import os

import onnx
import torch
import torch.nn as nn
from torchvision import models


def convert_to_onnx(pth_path, onnx_path, num_classes):
    # 加载PyTorch模型
    model = models.resnet50(weights=None)
    model.fc = nn.Linear(2048, num_classes)
    model.load_state_dict(torch.load(pth_path, weights_only=True))
    model.eval()

    # 创建示例输入
    dummy_input = torch.randn(1, 3, 224, 224)

    # 导出为ONNX格式
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        opset_version=11,
    )

    # 验证ONNX模型
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    print(f"模型已成功导出为 {onnx_path}")


if __name__ == "__main__":
    # 配置
    pth_path = "best_resnet50_model.pth"
    onnx_path = "model.onnx"
    num_classes = len(os.listdir("data/TrainingSet"))  # 获取类别数量

    # 转换模型
    convert_to_onnx(pth_path, onnx_path, num_classes)
